"""EDA预处理算法实现

提供多种EDA信号预处理算法：
- 基础预处理：异常值检测、平滑滤波、趋势去除
- 改进预处理：更精细的参数调优
- NeuroKit2预处理：集成NeuroKit2库的预处理功能

所有预处理函数返回统一格式的字典，包含原始数据、清洗数据、滤波数据等。
"""

import numpy as np
import pandas as pd
from scipy.signal import butter, filtfilt, savgol_filter, medfilt
from scipy.stats import zscore
from typing import Dict, Any, Optional
import logging

logger = logging.getLogger(__name__)

def preprocess_basic(eda_data: np.ndarray, sampling_rate: float = 4.0) -> Dict[str, Any]:
    """基础预处理算法
    
    执行基本的EDA信号预处理步骤：
    1. 异常值检测和替换（基于滑动窗口中位数）
    2. Savitzky-Golay平滑滤波
    3. 多项式趋势去除
    4. 低通滤波
    
    Args:
        eda_data: EDA信号数据
        sampling_rate: 采样率（Hz）
        
    Returns:
        包含各处理步骤结果的字典
    """
    try:
        logger.debug(f"开始基础预处理，数据长度: {len(eda_data)}")
        
        # 1. 异常值检测和替换
        window_size = max(8, int(sampling_rate * 2))  # 2秒窗口
        median_filtered = []
        
        for i in range(len(eda_data)):
            start_idx = max(0, i - window_size // 2)
            end_idx = min(len(eda_data), i + window_size // 2 + 1)
            window_data = eda_data[start_idx:end_idx]
            
            # 处理NaN值
            valid_window = window_data[~np.isnan(window_data)]
            if len(valid_window) > 0:
                median_filtered.append(np.median(valid_window))
            else:
                median_filtered.append(np.nan)
        
        median_filtered = np.array(median_filtered)
        
        # 计算与中位数的差异
        diff_from_median = np.abs(eda_data - median_filtered)
        valid_diff = diff_from_median[~np.isnan(diff_from_median)]
        
        if len(valid_diff) > 0:
            outlier_threshold = np.percentile(valid_diff, 95)
            outlier_mask = diff_from_median > outlier_threshold
        else:
            outlier_mask = np.zeros_like(eda_data, dtype=bool)
        
        # 替换异常值
        eda_cleaned = eda_data.copy()
        eda_cleaned[outlier_mask] = median_filtered[outlier_mask]
        
        # 2. 平滑滤波（Savitzky-Golay）
        window_length = min(9, len(eda_cleaned) // 4)
        if window_length % 2 == 0:
            window_length += 1
        window_length = max(3, window_length)
        
        # 处理NaN值
        valid_mask = ~np.isnan(eda_cleaned)
        if np.sum(valid_mask) > window_length:
            eda_smoothed = eda_cleaned.copy()
            valid_data = eda_cleaned[valid_mask]
            if len(valid_data) >= window_length:
                smoothed_valid = savgol_filter(valid_data, window_length, 3)
                eda_smoothed[valid_mask] = smoothed_valid
        else:
            eda_smoothed = eda_cleaned.copy()
        
        # 3. 趋势去除（二次多项式）
        x = np.arange(len(eda_smoothed))
        valid_indices = ~np.isnan(eda_smoothed)
        
        if np.sum(valid_indices) > 3:
            # 拟合多项式趋势
            trend_coeffs = np.polyfit(x[valid_indices], eda_smoothed[valid_indices], 2)
            trend = np.polyval(trend_coeffs, x)
            eda_detrended = eda_smoothed - trend
        else:
            eda_detrended = eda_smoothed.copy()
            trend = np.zeros_like(eda_smoothed)
        
        # 4. 低通滤波（Butterworth）
        try:
            nyquist = sampling_rate / 2
            cutoff_freq = min(0.5, nyquist * 0.8)  # 截止频率
            low = cutoff_freq / nyquist
            
            if low < 1.0:
                b, a = butter(4, low, btype='low')
                
                # 只对有效数据进行滤波
                if np.sum(valid_indices) > 10:
                    eda_filtered = eda_detrended.copy()
                    valid_data = eda_detrended[valid_indices]
                    filtered_valid = filtfilt(b, a, valid_data)
                    eda_filtered[valid_indices] = filtered_valid
                else:
                    eda_filtered = eda_detrended.copy()
            else:
                eda_filtered = eda_detrended.copy()
                
        except Exception as e:
            logger.warning(f"低通滤波失败: {str(e)}，使用去趋势数据")
            eda_filtered = eda_detrended.copy()
        
        # 统计信息
        processing_stats = {
            'outliers_detected': int(np.sum(outlier_mask)),
            'outlier_ratio': float(np.sum(outlier_mask) / len(eda_data)),
            'nan_count': int(np.sum(np.isnan(eda_data))),
            'valid_data_ratio': float(np.sum(valid_indices) / len(eda_data))
        }
        
        # 构建预处理结果字典，包含各个处理阶段的数据
        result = {
            'original': eda_data,           # 原始EDA数据，未经任何处理
            'cleaned': eda_cleaned,         # 清洁后的数据，已移除异常值和NaN
            'smoothed': eda_smoothed,       # 平滑处理后的数据，降低噪声
            'detrended': eda_detrended,     # 去趋势后的数据，移除长期趋势成分
            'filtered': eda_filtered,       # 滤波后的数据，应用低通滤波器
            'trend': trend,                 # 提取的趋势成分，用于分析长期变化
            'outlier_mask': outlier_mask,   # 异常值掩码，标记异常数据点位置
            'valid_mask': valid_indices,    # 有效数据掩码，标记可用数据点位置
            'processing_stats': processing_stats  # 处理统计信息，包含异常值数量等指标
        }
        
        logger.debug(f"基础预处理完成，异常值: {processing_stats['outliers_detected']}")
        return result
        
    except Exception as e:
        logger.error(f"基础预处理失败: {str(e)}")
        # 返回最小可用结果
        return {
            'original': eda_data,
            'cleaned': eda_data,
            'smoothed': eda_data,
            'detrended': eda_data,
            'filtered': eda_data,
            'outlier_mask': np.zeros_like(eda_data, dtype=bool),
            'valid_mask': ~np.isnan(eda_data),
            'processing_stats': {'error': str(e)}
        }

def preprocess_improved(eda_data: np.ndarray, sampling_rate: float = 4.0) -> Dict[str, Any]:
    """改进的预处理算法
    
    在基础预处理基础上增加：
    1. 更精细的异常值检测（Z-score + IQR）
    2. 自适应滤波参数
    3. 更好的边界处理
    4. 信号质量评估
    
    Args:
        eda_data: EDA信号数据
        sampling_rate: 采样率（Hz）
        
    Returns:
        包含各处理步骤结果的字典
    """
    try:
        logger.debug(f"开始改进预处理，数据长度: {len(eda_data)}")
        
        # 先执行基础预处理
        basic_result = preprocess_basic(eda_data, sampling_rate)
        
        # 1. 改进的异常值检测（Z-score + IQR方法）
        valid_data = eda_data[~np.isnan(eda_data)]
        if len(valid_data) > 10:
            # Z-score方法
            z_scores = np.abs(zscore(valid_data, nan_policy='omit'))
            z_outliers = z_scores > 3
            
            # IQR方法
            q1, q3 = np.percentile(valid_data, [25, 75])
            iqr = q3 - q1
            lower_bound = q1 - 1.5 * iqr
            upper_bound = q3 + 1.5 * iqr
            iqr_outliers = (valid_data < lower_bound) | (valid_data > upper_bound)
            
            # 合并异常值检测结果
            combined_outliers = z_outliers | iqr_outliers
            
            # 映射回原始数据
            improved_outlier_mask = np.zeros_like(eda_data, dtype=bool)
            valid_indices = ~np.isnan(eda_data)
            improved_outlier_mask[valid_indices] = combined_outliers
        else:
            improved_outlier_mask = basic_result['outlier_mask']
        
        # 2. 自适应中位数滤波
        # 根据信号变异性调整窗口大小
        signal_std = np.nanstd(eda_data)
        if signal_std > 2:
            med_window = max(3, int(sampling_rate * 1.5))  # 高变异性用较大窗口
        else:
            med_window = max(3, int(sampling_rate * 0.5))  # 低变异性用较小窗口
        
        if med_window % 2 == 0:
            med_window += 1
        
        eda_improved = eda_data.copy()
        if len(eda_data) >= med_window:
            # 对异常值位置应用中位数滤波
            median_filtered = medfilt(eda_data, kernel_size=med_window)
            eda_improved[improved_outlier_mask] = median_filtered[improved_outlier_mask]
        
        # 3. 自适应Savitzky-Golay滤波
        # 根据采样率和信号长度调整参数
        sg_window = min(
            max(5, int(sampling_rate * 2)),  # 至少2秒窗口
            len(eda_improved) // 3  # 不超过数据长度的1/3
        )
        if sg_window % 2 == 0:
            sg_window += 1
        
        valid_mask = ~np.isnan(eda_improved)
        if np.sum(valid_mask) >= sg_window:
            eda_sg_filtered = eda_improved.copy()
            valid_data = eda_improved[valid_mask]
            if len(valid_data) >= sg_window:
                smoothed_valid = savgol_filter(valid_data, sg_window, 3)
                eda_sg_filtered[valid_mask] = smoothed_valid
        else:
            eda_sg_filtered = eda_improved.copy()
        
        # 4. 改进的趋势去除（分段线性拟合）
        segment_length = max(60, int(len(eda_data) / 5))  # 分段长度
        eda_detrended_improved = eda_sg_filtered.copy()
        
        if len(eda_data) > segment_length * 2:
            for i in range(0, len(eda_data), segment_length):
                end_idx = min(i + segment_length, len(eda_data))
                segment = eda_sg_filtered[i:end_idx]
                segment_valid = ~np.isnan(segment)
                
                if np.sum(segment_valid) > 3:
                    x_seg = np.arange(len(segment))
                    trend_coeffs = np.polyfit(x_seg[segment_valid], segment[segment_valid], 1)
                    trend_seg = np.polyval(trend_coeffs, x_seg)
                    eda_detrended_improved[i:end_idx] = segment - trend_seg
        
        # 5. 双向滤波（更好的相位保持）
        try:
            nyquist = sampling_rate / 2
            cutoff_freq = min(0.4, nyquist * 0.7)  # 稍微保守的截止频率
            low = cutoff_freq / nyquist
            
            if low < 1.0 and np.sum(valid_mask) > 20:
                b, a = butter(3, low, btype='low')  # 降低阶数减少振荡
                eda_final = eda_detrended_improved.copy()
                valid_data = eda_detrended_improved[valid_mask]
                
                # 使用零相位滤波
                filtered_valid = filtfilt(b, a, valid_data)
                eda_final[valid_mask] = filtered_valid
            else:
                eda_final = eda_detrended_improved.copy()
                
        except Exception as e:
            logger.warning(f"改进滤波失败: {str(e)}")
            eda_final = eda_detrended_improved.copy()
        
        # 6. 信号质量评估
        quality_metrics = _assess_preprocessing_quality(eda_data, eda_final)
        
        # 更新结果
        result = basic_result.copy()
        result.update({
            'cleaned': eda_improved,
            'smoothed': eda_sg_filtered,
            'detrended': eda_detrended_improved,
            'filtered': eda_final,
            'outlier_mask': improved_outlier_mask,
            'quality_metrics': quality_metrics,
            'processing_stats': {
                **basic_result['processing_stats'],
                'improved_outliers': int(np.sum(improved_outlier_mask)),
                'sg_window_size': sg_window,
                'median_window_size': med_window
            }
        })
        
        logger.debug(f"改进预处理完成，质量评分: {quality_metrics.get('overall_quality', 'N/A')}")
        return result
        
    except Exception as e:
        logger.error(f"改进预处理失败: {str(e)}")
        # 回退到基础预处理
        return preprocess_basic(eda_data, sampling_rate)

def preprocess_neurokit2(eda_data: np.ndarray, sampling_rate: float = 4.0) -> Dict[str, Any]:
    """NeuroKit2预处理插件
    
    使用NeuroKit2库进行EDA预处理，如果库不可用则回退到改进算法。
    
    Args:
        eda_data: EDA信号数据
        sampling_rate: 采样率（Hz）
        
    Returns:
        包含各处理步骤结果的字典
    """
    try:
        import neurokit2 as nk
        logger.debug("使用NeuroKit2进行预处理")
        
        # 使用NeuroKit2处理
        signals, info = nk.eda_process(eda_data, sampling_rate=sampling_rate)
        
        # 提取处理结果
        eda_clean = signals['EDA_Clean'].values
        eda_phasic = signals.get('EDA_Phasic', eda_clean).values
        eda_tonic = signals.get('EDA_Tonic', np.zeros_like(eda_clean)).values
        
        # 计算有效性掩码
        valid_mask = ~np.isnan(eda_clean)
        
        result = {
            'original': eda_data,
            'cleaned': eda_clean,
            'smoothed': eda_clean,  # NeuroKit2已包含平滑
            'detrended': eda_phasic,  # 相位成分（去除基线）
            'filtered': eda_clean,
            'tonic': eda_tonic,  # 基线成分
            'phasic': eda_phasic,  # 相位成分
            'valid_mask': valid_mask,
            'outlier_mask': np.isnan(signals['EDA_Clean'].values),
            'neurokit2_info': info,
            'processing_stats': {
                'method': 'neurokit2',
                'valid_data_ratio': float(np.sum(valid_mask) / len(eda_data))
            }
        }
        
        logger.debug("NeuroKit2预处理完成")
        return result
        
    except ImportError:
        logger.info("NeuroKit2未安装，使用改进算法")
        return preprocess_improved(eda_data, sampling_rate)
    except Exception as e:
        logger.warning(f"NeuroKit2预处理失败: {str(e)}，使用改进算法")
        return preprocess_improved(eda_data, sampling_rate)

def _assess_preprocessing_quality(original: np.ndarray, processed: np.ndarray) -> Dict[str, float]:
    """评估预处理质量
    
    Args:
        original: 原始信号
        processed: 处理后信号
        
    Returns:
        质量指标字典
    """
    try:
        # 计算信噪比改善
        original_std = np.nanstd(original)
        processed_std = np.nanstd(processed)
        noise_reduction = original_std / processed_std if processed_std > 0 else 1
        
        # 计算相关性（保持信号特征）
        valid_mask = ~(np.isnan(original) | np.isnan(processed))
        if np.sum(valid_mask) > 10:
            correlation = np.corrcoef(original[valid_mask], processed[valid_mask])[0, 1]
        else:
            correlation = 0
        
        # 计算平滑度（梯度变化）
        if len(processed) > 1:
            gradient_original = np.abs(np.diff(original[~np.isnan(original)]))
            gradient_processed = np.abs(np.diff(processed[~np.isnan(processed)]))
            
            if len(gradient_original) > 0 and len(gradient_processed) > 0:
                smoothness_improvement = np.mean(gradient_original) / np.mean(gradient_processed)
            else:
                smoothness_improvement = 1
        else:
            smoothness_improvement = 1
        
        # 综合质量评分
        overall_quality = (
            min(noise_reduction, 5) * 0.3 +  # 噪声降低（限制在5倍以内）
            max(correlation, 0) * 0.4 +       # 信号保真度
            min(smoothness_improvement, 3) * 0.3  # 平滑度改善
        ) * 20  # 转换为0-100分
        
        return {
            'noise_reduction_ratio': float(noise_reduction),
            'signal_correlation': float(correlation),
            'smoothness_improvement': float(smoothness_improvement),
            'overall_quality': float(min(overall_quality, 100))
        }
        
    except Exception as e:
        logger.warning(f"质量评估失败: {str(e)}")
        return {
            'noise_reduction_ratio': 1.0,
            'signal_correlation': 0.0,
            'smoothness_improvement': 1.0,
            'overall_quality': 50.0
        }